# 1. 导入必备库（直接复制运行）
import pandas as pd
import numpy as np
# import matplotlib.pyplot as plt
from mlxtend.frequent_patterns import apriori
from mlxtend.frequent_patterns import association_rules

# 设置中文显示（避免乱码）

# 2. 数据加载与格式转换
# 加载数据集
df = pd.read_csv("supermarket_basket.csv")
# 查看数据基本信息
print(f"数据前5行：\n{df.head()}")
print()
print(f"数据形状（行数, 列数）：{df.shape}\n")
print(f"唯一商品列表：\n {df['product'].unique()}\n")
print(f"总订单数：{df['transaction_id'].nunique()}\n")

# 转换为“订单-商品”透视表（Apriori算法要求格式）
basket = df.pivot_table(
    index="transaction_id",  # 行：订单ID
    columns="product",  # 列：商品名称
    values="product",  # 用商品名称填充
    aggfunc="count",  # 计数（存在则计数≥1）
    fill_value=0  # 不存在填充0
).astype(bool)  # 转换为布尔型（True=购买，False=未购买）
print("透视表形状：", basket.shape)
print("透视表前5行：\n", basket.head())
print()
exit(0)

# 3. 挖掘频繁项集
# 最小支持度=0.05，挖掘频繁项集
frequent_itemsets = apriori(basket, min_support=0.05, use_colnames=True)
print("所有频繁项集数量：", len(frequent_itemsets))

# 筛选长度≥2的频繁项集，按支持度降序排序
frequent_itemsets_2plus = frequent_itemsets[frequent_itemsets["itemsets"].apply(len) >= 2]
frequent_itemsets_sorted = frequent_itemsets_2plus.sort_values("support", ascending=False)
print("长度≥2的频繁项集（前10条）：\n", frequent_itemsets_sorted.head(10))

# 4. 挖掘关联规则
# 基于频繁项集生成关联规则，最小置信度=0.3
rules = association_rules(frequent_itemsets, metric="confidence", min_threshold=0.3)
# 按提升度降序排序，显示核心列
rules_sorted = rules.sort_values("lift", ascending=False)[
    ["antecedents", "consequents", "support", "confidence", "lift"]]
print("关联规则（前10条）：\n", rules_sorted.head(10))

# 5. 结果可视化（Top5频繁二项集）
# 筛选频繁二项集（长度=2）
frequent_itemsets_2 = frequent_itemsets[frequent_itemsets["itemsets"].apply(len) == 2]
# 取Top5（支持度最高的5个二项集）
top5_itemsets = frequent_itemsets_2.sort_values("support", ascending=False).head(5)

# 处理标签（将frozenset转换为字符串，便于绘图）
top5_labels = [", ".join(list(itemset)) for itemset in top5_itemsets["itemsets"]]
top5_supports = top5_itemsets["support"].values

# 绘制柱状图

# 6. 强关联规则与运营建议
# 筛选提升度≥1.5的强关联规则
strong_rules = rules_sorted[rules_sorted["lift"] >= 1.5]
print("提升度≥1.5的强关联规则：\n", strong_rules)

# 运营策略建议（针对2条强关联规则，各写1条）
print("\n运营策略建议：")
# 示例（学生需根据实际挖掘结果调整）
rule1 = strong_rules.iloc[0]
print(f"1. 针对规则【{', '.join(list(rule1['antecedents']))} → {', '.join(list(rule1['consequents']))}】：")
print("   建议：__________（例：将A商品和B商品放在同一货架，提升连带购买率）")

rule2 = strong_rules.iloc[1]
print(f"2. 针对规则【{', '.join(list(rule2['antecedents']))} → {', '.join(list(rule2['consequents']))}】：")
print("   建议：__________（例：推出A+B商品组合优惠套餐，刺激消费）")
